#include <iostream>
#include <fstream>
#include <vector>
#include "utils.hpp"
#include "net.hpp"

int main(int main, char* argv[])
{
    const int numPoints = 1001;
    const float dt      = 2.0 / (numPoints - 1);

    // 准备数据
    std::vector<float> trueStat;
    std::vector<float> measuredVar;
    std::vector<float> controlVar;
    for (int i = 0; i < numPoints - 1; ++i)
    {
        const float t = double(i) * dt;
        // 控制变量（加速度）
        controlVar.push_back(acceleration(t));
        controlVar.push_back(acceleration(t));

        // 真实运动状态（解析计算结果）
        trueStat.push_back(position(t));
        trueStat.push_back(velocity(t));

        // 观测变量（包含假定的观测噪声）
        torch::Tensor tempPosi = torch::randn({1}) * 0.015 + position(t);
        measuredVar.push_back(tempPosi.item<float>());
        torch::Tensor tempVelo = torch::randn({1}) * 0.015 + velocity(t);
        measuredVar.push_back(tempVelo.item<float>());
    }

    // 初始化伸进网络（滤波器）及优化器
    std::vector<float> tempMatA{1, 0, dt, 1};
    std::vector<float> tempMatB{dt * dt / 2, 0, 0, dt};
    std::vector<float> tempMatH{1, 0, 0, 1};
    std::shared_ptr<Net> filter =
        std::make_shared<Net>(torch::from_blob(tempMatA.data(), {2, 2}, torch::kFloat),
                              torch::from_blob(tempMatB.data(), {2, 2}, torch::kFloat),
                              torch::from_blob(tempMatH.data(), {2, 2}, torch::kFloat), 2, 2);
    auto optimizer = std::make_shared<torch::optim::AdamW>(filter->parameters(), 0.1);

    // 开始训练，使用 0-1 秒的数据
    torch::Tensor target = torch::zeros({numPoints / 2, 2}, torch::kFloat);
    torch::Tensor index  = torch::zeros({numPoints / 2, 1}, torch::kLong);
    for (int i = 0; i < numPoints / 2; ++i)
    {

        target[i][0] = measuredVar[2 * i + 2];
        target[i][1] = measuredVar[2 * i + 3];

        index[i] = i;
    }

    auto dataSet = CustomDtaset(index, target).map(torch::data::transforms::Stack<>());

    auto dataLoader = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(
        std::move(dataSet), 32);

    float lossVal = 1.0, oldLoss = 1.0;
    filter->train();
    std::vector<float> res((numPoints / 2) * 2);
    while (std::abs(lossVal) > 1.e-6)
    {
        lossVal = 0.0;
        for (auto& batch : *dataLoader)
        {
            torch::Tensor idx = batch.data;
            torch::Tensor tag = batch.target;
            const int theSize = idx.size(0);

            torch::Tensor utm1Temp = torch::zeros({theSize, 2});
            torch::Tensor xtm1Temp = torch::zeros({theSize, 2});
            torch::Tensor ztTemp   = torch::zeros({theSize, 2});

            for (int i = 0; i < theSize; ++i)
            {
                const long theIdx = idx[i][0].item<long>();

                utm1Temp[i][0] = controlVar[2 * theIdx + 0];
                utm1Temp[i][1] = controlVar[2 * theIdx + 1];

                if (theIdx == 0)
                {
                    xtm1Temp[i][0] = 0;
                    xtm1Temp[i][1] = 0;
                }
                else
                {
                    xtm1Temp[i][0] = res[2 * (theIdx - 1) + 0];
                    xtm1Temp[i][1] = res[2 * (theIdx - 1) + 1];
                }

                ztTemp[i][0] = measuredVar[2 * theIdx + 2];
                ztTemp[i][1] = measuredVar[2 * theIdx + 3];
            }
            //
            torch::Tensor out = filter->forward(utm1Temp, xtm1Temp, ztTemp);

            //
            std::vector<float> aveValOut(2, 0.0), aveValTag(2, 0.0);
            for (int i = 0; i < theSize; ++i)
            {
                aveValOut[0] += out[i][0].item<float>();
                aveValOut[1] += out[i][1].item<float>();
                aveValTag[0] += tag[i][0].item<float>();
                aveValTag[1] += tag[i][1].item<float>();
            }
            aveValOut[0] /= theSize;
            aveValOut[1] /= theSize;
            aveValTag[0] /= theSize;
            aveValTag[1] /= theSize;
            torch::Tensor aveOut = torch::zeros({theSize, 2}, torch::kFloat);
            torch::Tensor aveTag = torch::zeros({theSize, 2}, torch::kFloat);
            for (int i = 0; i < theSize; ++i)
            {
                aveOut[i][0] = aveValOut[0];
                aveOut[i][1] = aveValOut[1];
                aveTag[i][0] = aveValTag[0];
                aveTag[i][1] = aveValTag[1];
            }
            torch::Tensor loss = torch::var(out - aveOut) + torch::var(tag - aveOut);

            //
            optimizer->zero_grad();
            loss.backward();
            optimizer->step();

            //
            lossVal += loss.item<float>();

            for (int i = 0; i < theSize; ++i)
            {
                const long theIdx = idx[i][0].item<long>();

                res[2 * theIdx + 0] = out[i][0].item<float>();
                res[2 * theIdx + 1] = out[i][1].item<float>();
            }
        }

        if (std::abs(lossVal / oldLoss - 1.0) < 1.e-5) break;
        std::cout << "Train Loss : " << lossVal << std::endl;
        oldLoss = lossVal;
    }

    // 输出结果
    filter->eval();
    std::ofstream os("data_from_train.csv");
    torch::Tensor utm1Temp = torch::zeros({1, 2});
    torch::Tensor xtm1Temp = torch::zeros({1, 2});
    torch::Tensor ztTemp   = torch::zeros({1, 2});
    for (int i = 0; i < numPoints - 2; ++i)
    {
        const float t = double(i) * dt;

        utm1Temp[0][0] = controlVar[2 * i + 0];
        utm1Temp[0][1] = controlVar[2 * i + 1];
        if (i == 0)
        {
            xtm1Temp[0][0] = 0;
            xtm1Temp[0][1] = 0;
        }
        ztTemp[0][0] = measuredVar[2 * i + 2];
        ztTemp[0][1] = measuredVar[2 * i + 3];

        auto out = filter->forward(utm1Temp, xtm1Temp, ztTemp);

        os << t << "," << controlVar[2 * i + 0] << "," << controlVar[2 * i + 1] << ","
           << trueStat[2 * i + 0] << "," << trueStat[2 * i + 1] << "," << measuredVar[2 * i + 2]
           << "," << measuredVar[2 * i + 3] << "," << out[0][0].item<float>() << ","
           << out[0][1].item<float>() << std::endl;
        xtm1Temp[0][0] = out[0][0].item<float>();
        xtm1Temp[0][1] = out[0][1].item<float>();
    }
    os.close();
    return 0;
}
